"""
速率限制器模块
防止API滥用，保护系统资源
支持基于IP的请求频率限制
"""

import time
import logging
from typing import Dict, List
from fastapi import HTTPException, Request
from .config import settings

logger = logging.getLogger(__name__)

# 存储每个IP的请求记录
# 格式: {ip_address: [timestamp1, timestamp2, ...]}
REQUEST_RECORDS: Dict[str, List[float]] = {}


def get_rate_limit_config() -> tuple:
    """
    获取速率限制配置
    
    Returns:
        tuple: (请求数限制, 时间窗口)
    """
    return settings.rate_limit_requests, settings.rate_limit_window


def get_client_ip(request: Request) -> str:
    """
    获取客户端真实IP地址
    优先从代理头中获取，避免代理服务器影响
    
    Args:
        request: FastAPI请求对象
        
    Returns:
        str: 客户端IP地址
    """
    # 优先从X-Forwarded-For头获取真实IP
    forwarded = request.headers.get("X-Forwarded-For")
    if forwarded:
        # X-Forwarded-For可能包含多个IP，取第一个
        return forwarded.split(",")[0].strip()
    
    # 从X-Real-IP头获取
    real_ip = request.headers.get("X-Real-IP")
    if real_ip:
        return real_ip.strip()
    
    # 最后使用直连IP
    return request.client.host if request.client else "unknown"


def check_rate_limit(client_ip: str, limit: int = None, window: int = None):
    """
    检查指定IP的请求频率是否超限
    如果超限则抛出HTTPException
    
    Args:
        client_ip: 客户端IP地址
        limit: 请求限制数量，默认使用配置值
        window: 时间窗口（秒），默认使用配置值
        
    Raises:
        HTTPException: 当请求频率超限时抛出429错误
    """
    current_time = time.time()
    limit = limit or settings.rate_limit_requests
    window = window or settings.rate_limit_window
    
    # 获取该IP的请求记录
    if client_ip not in REQUEST_RECORDS:
        REQUEST_RECORDS[client_ip] = []
    
    requests = REQUEST_RECORDS[client_ip]
    
    # 清理过期的请求记录（超出时间窗口的请求）
    cutoff_time = current_time - window
    requests[:] = [req_time for req_time in requests if req_time > cutoff_time]
    
    # 检查是否超出限制
    if len(requests) >= limit:
        # 计算还需要等待多长时间
        oldest_request = min(requests)
        retry_after = int(window - (current_time - oldest_request)) + 1
        
        logger.warning(
            f"速率限制触发 - IP: {client_ip}, "
            f"请求数: {len(requests)}/{limit}, "
            f"时间窗口: {window}秒, "
            f"需等待: {retry_after}秒"
        )
        
        raise HTTPException(
            status_code=429,
            detail={
                "message": f"请求过于频繁，每{window}秒最多允许{limit}次请求",
                "retry_after": retry_after,
                "current_requests": len(requests),
                "limit": limit,
                "window": window
            },
            headers={"Retry-After": str(retry_after)}
        )
    
    # 记录本次请求
    requests.append(current_time)
    
    # 定期清理REQUEST_RECORDS，避免内存泄漏
    _cleanup_old_records(current_time, window)


def _cleanup_old_records(current_time: float, window: int):
    """
    清理过期的请求记录，防止内存泄漏
    
    Args:
        current_time: 当前时间戳
        window: 时间窗口
    """
    # 每100次请求清理一次
    if len(REQUEST_RECORDS) > 100 and int(current_time) % 10 == 0:
        cutoff_time = current_time - window * 2  # 保留2倍窗口时间的数据
        
        # 删除完全过期的IP记录
        expired_ips = []
        for ip, requests in REQUEST_RECORDS.items():
            if not requests or max(requests) < cutoff_time:
                expired_ips.append(ip)
        
        for ip in expired_ips:
            del REQUEST_RECORDS[ip]
        
        if expired_ips:
            logger.debug(f"清理了{len(expired_ips)}个过期IP的请求记录")


def get_rate_limit_status(client_ip: str) -> Dict[str, any]:
    """
    获取指定IP的速率限制状态
    
    Args:
        client_ip: 客户端IP地址
        
    Returns:
        Dict: 包含当前请求数、限制等信息的字典
    """
    current_time = time.time()
    cutoff_time = current_time - settings.rate_limit_window
    
    if client_ip not in REQUEST_RECORDS:
        current_requests = 0
    else:
        # 计算时间窗口内的请求数
        requests = REQUEST_RECORDS[client_ip]
        current_requests = len([req for req in requests if req > cutoff_time])
    
    remaining = max(0, settings.rate_limit_requests - current_requests)
    
    return {
        "ip": client_ip,
        "current_requests": current_requests,
        "limit": settings.rate_limit_requests,
        "window": settings.rate_limit_window,
        "remaining": remaining,
        "reset_time": int(current_time + settings.rate_limit_window)
    }


def reset_rate_limit(client_ip: str = None):
    """
    重置速率限制记录
    
    Args:
        client_ip: 要重置的IP地址，如果为None则重置所有记录
    """
    if client_ip:
        REQUEST_RECORDS.pop(client_ip, None)
        logger.info(f"已重置IP {client_ip} 的速率限制记录")
    else:
        REQUEST_RECORDS.clear()
        logger.info("已重置所有速率限制记录") 